# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import functools
import io
import json

import tornado.gen
import requests

from coin.base.timestamp import get_timestamp
import coin.flow.subscriber as flow_subscriber
import coin.flow.topic_map as topic_map

from coin.exchange.base.rest.task_scheduler import Scheduler, wrap_async_task
from coin.exchange.okex_futures_swap.kr_rest.native_private_client import OkexSwapNativePrivateClient
from coin.exchange.okex_futures_swap.kr_rest.private_client import OkexFuturesSwapPrivateParser
from coin.exchange.okex_futures_swap.ws.private_subscriber import OkexSwapWsPrivateClient
from coin.exchange.okex_v3.kr_rest.native_private_client import OkexFuturesOrderStatus
from coin.exchange.base.order_gateway_util import (
    PrivateExchangeMessage,
    convert_to_requests_response,
)
from coin.exchange.okex_futures_v3.order_gateway import (
    OkexFuturesOrderGatewayBase,
    OkexFuturesOrderGatewayInfo,
    OkexCurrency,
    OkexFuturesProduct,
    print_working_orders,
    OkexOrderType,
    OkexOrderProperty,
    OrderType,
)


class OkexFuturesSwapOrderGatewayInfo(OkexFuturesOrderGatewayInfo):
  def update_balance(self, msg):
    account_balance = OkexFuturesSwapPrivateParser.parse_balance(msg)
    for balance in account_balance.each_balance:
      currency = OkexCurrency.FromStr(balance.currency)
      self.set_balance(currency, balance.total, balance)

    for balance_info in msg['info']:
      product = OkexFuturesProduct.FromStrNativeProduct(balance_info['instrument_id'])
      currency = product.base
      self._unrealized_pnl[currency.currency] = float(balance_info['unrealized_pnl'])

  def update_position(self, message):
    account_position = OkexFuturesSwapPrivateParser.parse_position(message)
    for position in account_position.each_position:
      product = OkexFuturesProduct.FromStrProduct(position.symbol)
      self.set_position(product, position.net_position, position)

  def update_ws_position(self, message):
    position = OkexFuturesSwapPrivateParser.parse_ws_product_position(message)
    product = OkexFuturesProduct.FromStrProduct(position.symbol)
    self.set_position(product, position.net_position, position)


class OkexFuturesSwapOrderGateway(OkexFuturesOrderGatewayBase):
  _exchange = 'Okex'
  _market_type = 'Futures'
  _api_version = 'v3_swap'

  OGInfoType = OkexFuturesSwapOrderGatewayInfo

  def start(self, queue, ioloop, flow_sub=None):
    api_host = self._og_config.rest_root_url \
      if self._og_config.HasField('rest_root_url') else None
    self._private_client = OkexSwapNativePrivateClient(key_file=self._config.key_file,
                                                       api_host=api_host,
                                                       use_async=True,
                                                       rest_lib=self._og_config.rest_lib)

    topic_string = 'okex_swap_rest_private_1'
    topic_id = topic_map.search_topic_id(topic_string)[0][1]
    self._ioloop = ioloop
    self._publisher = queue.get_writer(topic_id=topic_id)
    self._scheduler = Scheduler(ioloop=ioloop, logger=self._logger)
    self._prepare_private_subscriber()
    self._install_periodic_query()

    # Subscribe to flow for feed.
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe(topic_string, self.on_private_exchange_msg)

  def _prepare_private_subscriber(self):
    if self._private_subscriber is not None:
      self._private_subscriber.stop()

    def on_ready():
      products = [prod.native_symbol_v3 for prod in self._og_info.products]
      self._private_subscriber.prepare_swap_subscription(products)

    self._private_subscriber = OkexSwapWsPrivateClient(
        key_file=self._config.key_file,
        ioloop=self._ioloop,
        notification_callback=self._notification_callback,
        close_callback=functools.partial(self._prepare_private_subscriber),
        on_ready=on_ready,
    )
    self._private_subscriber.start()

  def _async_submit_order(self, order, post_only=None):
    assert order.order_id is not None, order
    assert order.order_type == OrderType.LIMIT, order

    if post_only:
      post_only = OkexOrderProperty.POST_ONLY
    else:
      post_only = None

    formatter = self._order_formatter[order.product]
    params = {
        'instrument_id': order.product.native_symbol_v3,
        'price': formatter.format_price(order.price),
        'size': formatter.format_qty(order.qty),
        'order_side': order.order_side,
        'order_type': OkexOrderType(order.order_type.name.lower()),
        'order_property': post_only,
    }
    fut = self._private_client.place_order(**params)
    return fut

  @tornado.gen.coroutine
  def _query_open_orders(self, instrument_id):
    request_timestamp = get_timestamp()
    # SWAP does not support query open orders
    response = yield self._private_client.query_order_list(instrument_id=instrument_id,
                                                           status=OkexFuturesOrderStatus.NEW)
    response = convert_to_requests_response(response)
    if response.status_code != 200:
      self._logger.error('Query new order failed!')
      return
    new_orders = response.json()

    response = yield self._private_client.query_order_list(
        instrument_id=instrument_id, status=OkexFuturesOrderStatus.PART_FILLED)
    response = convert_to_requests_response(response)
    if response.status_code != 200:
      self._logger.error('Query partially filled order failed!')
      return

    partially_filled_orders = response.json()
    open_orders = {'order_info': new_orders['order_info'] + partially_filled_orders['order_info']}

    response = requests.Response()
    response.status_code = 200
    response.encoding = 'UTF-8'
    response.raw = io.BytesIO(json.dumps(open_orders).encode())
    timestamp = get_timestamp()
    msg = PrivateExchangeMessage(
        msg_type=self._MsgType.REST_OPEN_ORDERS,
        data=response,
        received_timestamp=timestamp,
        request_params={'instrument_id': instrument_id},
        request_timestamp=request_timestamp,
    )
    self._publisher.write(timestamp, msg)

  @tornado.gen.coroutine
  def _query_done_orders(self, instrument_id):
    request_timestamp = get_timestamp()
    # SWAP does not support query done orders
    response = yield self._private_client.query_order_list(instrument_id=instrument_id,
                                                           status=OkexFuturesOrderStatus.CANCELED)
    response = convert_to_requests_response(response)
    if response.status_code != 200:
      self._logger.error('Query cancel order failed!')
      return
    cancel_orders = response.json()

    response = yield self._private_client.query_order_list(instrument_id=instrument_id,
                                                           status=OkexFuturesOrderStatus.FILLED)
    response = convert_to_requests_response(response)
    if response.status_code != 200:
      self._logger.error('Query filled order failed!')
      return
    filled_orders = response.json()

    done_orders = {'order_info': cancel_orders['order_info'] + filled_orders['order_info']}

    response = requests.Response()
    response.status_code = 200
    response.encoding = 'UTF-8'
    response.raw = io.BytesIO(json.dumps(done_orders).encode())
    timestamp = get_timestamp()
    msg = PrivateExchangeMessage(
        msg_type=self._MsgType.REST_DONE_ORDERS,
        data=response,
        received_timestamp=timestamp,
        request_params={'instrument_id': instrument_id},
        request_timestamp=request_timestamp,
    )
    self._publisher.write(timestamp, msg)

  @tornado.gen.coroutine
  def _query_account_balance(self):
    # emulate a single call of account balance query
    request_timestamp = get_timestamp()
    aggregated = {'result': True, 'info': []}
    for product in self._og_info.products:
      response = yield self._private_client.query_single_account(
          instrument_id=product.native_symbol_v3)
      response = convert_to_requests_response(response)
      if response.status_code != 200:
        self._logger.error('query %s balance failed!' % product)
        continue
      balance_info = response.json()
      aggregated['info'].append(balance_info['info'])

    response = requests.Response()
    response.url = ('query_account_balance (%s, %s, %s)' %
                    (self._exchange, self._market_type, self._api_version))
    response.status_code = 200
    response.encoding = 'UTF-8'
    response.raw = io.BytesIO(json.dumps(aggregated).encode())
    timestamp = get_timestamp()
    msg = PrivateExchangeMessage(
        msg_type=self._MsgType.REST_BALANCE,
        data=response,
        received_timestamp=timestamp,
        request_params=None,
        request_timestamp=request_timestamp,
    )
    self._publisher.write(timestamp, msg)

  @tornado.gen.coroutine
  def _install_periodic_query(self):
    assert self._scheduler is not None, 'self._scheduler is not set!'
    assert self._publisher is not None, 'self._publisher is not set!'

    # query account balance
    self._scheduler.add_periodic_task(period=self._og_config.account_query_period,
                                      job_func=self._query_account_balance)

    for product in self._og_info.products:
      instrument_id = product.native_symbol_v3
      # query active orders
      yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
      self._scheduler.add_periodic_task(period=self._og_config.orders_snapshot_query_period,
                                        job_func=self._query_open_orders,
                                        instrument_id=instrument_id)

      # query completed orders (filled + cancelled)
      yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
      self._scheduler.add_periodic_task(period=self._og_config.fill_query_period,
                                        job_func=self._query_done_orders,
                                        instrument_id=instrument_id)

      # query positions
      yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
      query_position_task = wrap_async_task(
          self._private_client.query_contract_positions,
          msg_type=self._MsgType.REST_POSITION,
          publisher=self._publisher,
          func_params=dict(instrument_id=instrument_id),
      )
      self._scheduler.add_periodic_task(
          period=self._og_config.user_query_period,
          job_func=query_position_task,
      )

    # Periodically print working order list.
    print_info = functools.partial(print_working_orders, self._logger, self)
    self._scheduler.add_periodic_task(period=10, job_func=print_info)
    self._scheduler.start()
